from mem0 import Memory
import os
import sys
from mem0.configs.base import MemoryConfig
from mem0.llms.configs import LlmConfig

# 添加项目根目录到Python路径，以便导入config_reader模块
current_dir = os.path.dirname(os.path.abspath(__file__))
project_root = os.path.dirname(os.path.dirname(current_dir))
sys.path.append(project_root)

from config_reader import get_openai_hk_api_key

# 从配置文件获取OpenAI HK API密钥
api_key = get_openai_hk_api_key()
print(f"从配置文件读取的 OPENAI_HK_API_KEY: {api_key[:20]}...")

# 配置API网关 - 使用指定的主网关和备用网关
api_base = "https://api.openai-hk.com/v1"  # 主网关

# 创建LLM配置
llm_config = LlmConfig(
    provider="openai",
    config={
        "api_key": api_key,  # 你的OpenAI API密钥
        "openai_base_url": api_base,  # API基础地址，默认即可
        # "model": "gpt-3.5-turbo",  # 指定使用的模型
        # "temperature": 0.7  # 其他模型参数
    }
)

# 创建内存配置
config = MemoryConfig(llm=llm_config)


# 初始化记忆模块，指定自定义API基础地址
m = Memory(config=config)

# 为用户添加记忆
messages = [
    {
        "role": "user",
        "content": "I like to drink coffee in the morning and go for a walk"
    }
]
result = m.add(messages, user_id="alice", metadata={"category": "preferences"})

# 搜索相关记忆
related_memories = m.search("Should I drink coffee or tea?", user_id="alice")

print("相关记忆:", related_memories)
